#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
crop_tif.py

根据经纬度边界框裁剪 GeoTIFF 文件，支持单文件和批量处理。

功能：
- 根据提供的经纬度四个数值（minx, miny, maxx, maxy）裁剪 TIF
- 支持单文件或目录批量处理
- 自动处理不同 CRS 的影像
- 保持原始分辨率和数据类型
- 支持指定输出 CRS
- LZW 压缩输出

依赖：rasterio, numpy

用法示例：
1) 单文件裁剪：
   python crop_tif.py --input input.tif --output cropped.tif --bounds 120.0 30.0 130.0 40.0

2) 批量裁剪：
   python crop_tif.py --input-dir ./tiles --output-dir ./cropped --bounds 120.0 30.0 130.0 40.0

3) 指定输出 CRS：
   python crop_tif.py --input input.tif --output cropped.tif --bounds 120.0 30.0 130.0 40.0 --dst-crs EPSG:4326
"""

import argparse
import os
from typing import List, Optional, Tuple

import numpy as np
import rasterio
from rasterio.warp import calculate_default_transform, reproject, Resampling
from rasterio.windows import from_bounds


class TifCropper:
    """TIF 裁剪器"""
    
    def __init__(self, bounds: Tuple[float, float, float, float], dst_crs: Optional[str] = None):
        """
        初始化裁剪器
        
        Args:
            bounds: 裁剪边界框 (minx, miny, maxx, maxy)
            dst_crs: 输出 CRS，如果为 None 则使用输入文件的 CRS
        """
        self.bounds = bounds
        self.dst_crs = dst_crs
        self.minx, self.miny, self.maxx, self.maxy = bounds
        
        # 验证边界框
        if self.minx >= self.maxx or self.miny >= self.maxy:
            raise ValueError("无效的边界框：minx >= maxx 或 miny >= maxy")
    
    def crop_single_file(self, input_path: str, output_path: str) -> str:
        """
        裁剪单个 TIF 文件
        
        Args:
            input_path: 输入文件路径
            output_path: 输出文件路径
            
        Returns:
            输出文件路径
        """
        with rasterio.open(input_path) as src:
            # 获取源文件的 CRS
            src_crs = src.crs
            
            # 确定输出 CRS
            output_crs = self.dst_crs if self.dst_crs else src_crs
            
            if src_crs != output_crs:
                # 需要重投影
                # 将边界框从输出 CRS 转换到源 CRS
                from rasterio.warp import transform
                src_bounds = transform(
                    output_crs, src_crs,
                    [self.minx, self.maxx], [self.miny, self.maxy]
                )
                src_minx, src_maxx = src_bounds[0]
                src_miny, src_maxy = src_bounds[1]
                
                # 确保边界框在源文件范围内
                src_minx = max(src_minx, src.bounds.left)
                src_miny = max(src_miny, src.bounds.bottom)
                src_maxx = min(src_maxx, src.bounds.right)
                src_maxy = min(src_maxy, src.bounds.top)
                
                if src_minx >= src_maxx or src_miny >= src_maxy:
                    raise ValueError(f"边界框与文件 {input_path} 无重叠区域")
                
                # 创建裁剪窗口
                window = from_bounds(src_minx, src_miny, src_maxx, src_maxy, src.transform)
                
                # 读取数据
                data = src.read(window=window)
                src_transform = src.window_transform(window)
                
                # 计算目标变换
                dst_transform, dst_width, dst_height = calculate_default_transform(
                    src_crs, output_crs,
                    data.shape[1], data.shape[2],
                    src_minx, src_miny, src_maxx, src_maxy
                )
                
                # 重投影数据
                dst_data = np.empty((data.shape[0], dst_height, dst_width), dtype=data.dtype)
                reproject(
                    data, dst_data,
                    src_transform=src_transform,
                    dst_transform=dst_transform,
                    src_crs=src_crs,
                    dst_crs=output_crs,
                    resampling=Resampling.nearest
                )
                
                # 更新元数据
                profile = src.profile.copy()
                profile.update({
                    'crs': output_crs,
                    'transform': dst_transform,
                    'width': dst_width,
                    'height': dst_height,
                    'compress': 'lzw'
                })
                
            else:
                # 不需要重投影，直接裁剪
                # 确保边界框在文件范围内
                crop_minx = max(self.minx, src.bounds.left)
                crop_miny = max(self.miny, src.bounds.bottom)
                crop_maxx = min(self.maxx, src.bounds.right)
                crop_maxy = min(self.maxy, src.bounds.top)
                
                if crop_minx >= crop_maxx or crop_miny >= crop_maxy:
                    raise ValueError(f"边界框与文件 {input_path} 无重叠区域")
                
                # 创建裁剪窗口
                window = from_bounds(crop_minx, crop_miny, crop_maxx, crop_maxy, src.transform)
                
                # 读取数据
                data = src.read(window=window)
                dst_transform = src.window_transform(window)
                
                # 更新元数据
                profile = src.profile.copy()
                profile.update({
                    'transform': dst_transform,
                    'width': data.shape[2],
                    'height': data.shape[1],
                    'compress': 'lzw'
                })
            
            # 创建输出目录
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 写入文件
            with rasterio.open(output_path, 'w', **profile) as dst:
                for i in range(data.shape[0]):
                    dst.write(data[i], i + 1)
        
        return output_path
    
    def crop_directory(self, input_dir: str, output_dir: str) -> List[str]:
        """
        批量裁剪目录中的 TIF 文件
        
        Args:
            input_dir: 输入目录
            output_dir: 输出目录
            
        Returns:
            输出文件路径列表
        """
        # 查找所有 TIF 文件
        tif_files = []
        for root, dirs, files in os.walk(input_dir):
            for file in files:
                if file.lower().endswith(('.tif', '.tiff')):
                    tif_files.append(os.path.join(root, file))
        
        if not tif_files:
            print(f"在目录 {input_dir} 中未找到 TIF 文件")
            return []
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        output_files = []
        for input_path in tif_files:
            # 生成输出文件名
            rel_path = os.path.relpath(input_path, input_dir)
            output_path = os.path.join(output_dir, rel_path)
            
            try:
                self.crop_single_file(input_path, output_path)
                output_files.append(output_path)
                print(f"裁剪完成: {os.path.basename(input_path)} -> {os.path.basename(output_path)}")
            except Exception as e:
                print(f"裁剪失败 {os.path.basename(input_path)}: {e}")
        
        return output_files


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='根据经纬度边界框裁剪 GeoTIFF 文件')
    
    # 输入参数组
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument('--input', help='输入 TIF 文件路径')
    input_group.add_argument('--input-dir', help='输入目录（递归搜索 .tif/.tiff）')
    
    # 输出参数
    parser.add_argument('--output', help='输出 TIF 文件路径（单文件模式必填）')
    parser.add_argument('--output-dir', help='输出目录（批量模式必填）')
    
    # 裁剪参数
    parser.add_argument('--bounds', nargs=4, type=float, required=True,
                       metavar=('MINX', 'MINY', 'MAXX', 'MAXY'),
                       help='裁剪边界框 (minx, miny, maxx, maxy)')
    parser.add_argument('--dst-crs', help='输出 CRS（如 EPSG:4326），默认使用输入文件的 CRS')
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    # 验证输入
    if args.input and not os.path.exists(args.input):
        raise SystemExit(f"输入文件不存在: {args.input}")
    if args.input_dir and not os.path.exists(args.input_dir):
        raise SystemExit(f"输入目录不存在: {args.input_dir}")
    
    # 创建裁剪器
    cropper = TifCropper(bounds=tuple(args.bounds), dst_crs=args.dst_crs)
    
    try:
        if args.input:
            # 单文件模式
            if not args.output:
                raise SystemExit("单文件模式需要提供 --output")
            
            print(f"正在裁剪文件: {args.input}")
            output_path = cropper.crop_single_file(args.input, args.output)
            print(f"裁剪完成: {output_path}")
            
        else:
            # 批量模式
            if not args.output_dir:
                raise SystemExit("批量模式需要提供 --output-dir")
            
            print(f"正在批量裁剪目录: {args.input_dir}")
            output_files = cropper.crop_directory(args.input_dir, args.output_dir)
            print(f"批量裁剪完成，共处理 {len(output_files)} 个文件")
            
    except Exception as e:
        print(f"裁剪过程中发生错误: {e}")
        raise


if __name__ == '__main__':
    main()

